import sys
import math
import numpy
import random
from SimPy.Simulation import *

def not_in(a, b, pop=False, elems=set()):
    '''nodes from set a that are not in set b'''
    for elem in a:
        if elem not in b:
            if pop: elems.add(elem)
            yield elem
    if pop:
        for elem in elems:
            a.remove(elem)
        elems.clear()

class Expon:
    @staticmethod
    def cdf(lambd, x):
        return 1 - math.exp(-lambd * x)

    @staticmethod
    def rvs(lambd):
        return -math.log(1 - random.random())/(lambd)

class Node(Process):
    basetime = 3600*24
    availability = 0.66
    abandon_probability = 0.5
    # Time Required to upload a block.
    block_time = 0.0003 # 1 minute approx
    
    online_nodes = set()
    offline_nodes = set()
    departed_nodes = set()
    
    id = 0
    node_set_changed = None

    def __init__(self, *args, **kargs):
        Process.__init__(self, *args, **kargs)
        self.set_id()
        self.time_on = Node.availability
        self.time_off = (1-Node.availability)
        self.lambd_on = 1./self.time_on
        self.lambd_off = 1./self.time_off

    def set_id(self):
        self.id = Node.id
        Node.id += 1

    def update_next_event(self, period):
        next_time = self.sim.now() + period
        if Node.next_event < self.sim.now() or \
           Node.next_event > next_time or \
           Node.next_event == None:
            Node.next_event = next_time

    def run(self):
        while True:
            Node.online_nodes.add(self)
            Node.node_set_changed.signal()
            yield hold, self, Expon.rvs(self.lambd_on)*Node.basetime
            Node.online_nodes.remove(self)

            if random.random()<=Node.abandon_probability:
                Node.node_set_changed.signal()
                #Node.departed_nodes.add(self)
                break
            
            Node.offline_nodes.add(self)
            Node.node_set_changed.signal()
            yield hold, self, Expon.rvs(self.lambd_off)*Node.basetime
            Node.offline_nodes.remove(self)

class RepairTimeout(Process):
    def __init__(self, id, *args, **kargs):
        Process.__init__(self, *args, **kargs)
        self.timeout = 1
        self.event = SimEvent(sim=self.sim, name='repair_timeout_'+str(id))

    def set(self, t):
        self.timeout = t

    def run(self):
        while True:
            yield hold, self, self.timeout
            self.event.signal()

class Repair(Process):
    required_blocks = 50
    blocks_per_node = 1
    active_repairs = set()
    id = 0

    def __init__(self, *args, **kargs):
        Process.__init__(self, *args, **kargs)
        self.id = Repair.id
        Repair.id += 1

    def finished(self):
        n = sum(self.downloaded_blocks.itervalues())
        return n>=Repair.required_blocks

    def run(self):
        Repair.active_repairs.add(self)
        
        timeout = RepairTimeout(self.id, sim=self.sim)
        self.sim.activate(timeout, timeout.run())

        self.downloading_node = random.choice(list(Node.online_nodes))
        self.downloaded_blocks = {}
        self.downloaded_blocks[self.downloading_node] = 0
        self.downloaded_partial = 0

        start_time = self.sim.now()
        while not self.finished():
            time = self.sim.now()
            block_time = Node.block_time*Node.basetime
            timeout.set(block_time)
            yield waitevent, self, (Node.node_set_changed, timeout.event)
            time = self.sim.now() - time
            if block_time>0:
                partial = time/block_time

                # Append partial downloaded information and update downloaded blocks
                self.downloaded_partial += partial
                if self.downloaded_partial > 1:
                    n = int(self.downloaded_partial)
                    self.downloaded_blocks[self.downloading_node] += n
                    self.downloaded_partial -= n

                # Disconnected node => partial information is lost.
                if self.downloading_node not in Node.online_nodes:
                    self.downloaded_partial = 0
                    self.downloading_node = random.choice(list(Node.online_nodes))
                    if not self.downloaded_blocks.has_key(self.downloading_node):
                        self.downloaded_blocks[self.downloading_node] = 0
        
        self.cancel(timeout)
        node = Node(sim=self.sim)
        self.sim.activate(node, node.run())
        print self.sim.now(), 'endrepair', self.sim.now()-start_time
        Repair.active_repairs.remove(self)

class Oracle(Process):
    def __init__(self, observe, *args, **kargs):
        Process.__init__(self, *args, **kargs)
        self.observe = observe

    def run(self):
        while True:
            yield waitevent, self, Node.node_set_changed
            if len(Node.online_nodes)<self.observe:
                repair = Repair(sim=self.sim)
                self.sim.activate(repair, repair.run())
                print self.sim.now(), 'oracle'

class StorageSystem(Process):
    def __init__(self, targeted_nodes, repair_rate, threshold, required, *args, **kargs):
        Process.__init__(self, *args, **kargs)
        
        self.targeted_nodes = targeted_nodes
        self.threshold = threshold
        self.repair_rate = repair_rate
        
        Repair.required_blocks = required
        
        for i in xrange(self.targeted_nodes):
            node = Node(sim=self.sim)
            self.sim.activate(node, node.run())

        Node.node_set_changed = SimEvent(sim=self.sim, name='node_set_changed')
        self.oracle = Oracle(self.threshold, sim=self.sim)

    def run(self):
        self.sim.activate(self.oracle, self.oracle.run())
        tolerance = int(numpy.ceil(Repair.required_blocks/Repair.blocks_per_node))
        while len(Node.online_nodes)+len(Node.offline_nodes)>=tolerance:
            repair = Repair(sim=self.sim)
            self.sim.activate(repair, repair.run())
            print self.sim.now(), 'inirepair', len(Node.online_nodes), len(Node.offline_nodes)
            yield hold, self, Node.basetime/self.repair_rate

        print self.sim.now(), 'datalost'
        self.cancel(self.oracle)
            
def main(type):
    n = 100
    required = 70*Repair.blocks_per_node
    threshold = 75

    a = Node.abandon_probability
    m = Node.availability
    repair_rate_avail = a*n/m
    repair_rate_stored = a*n/(a*m-a+1)
    f = repair_rate_avail/repair_rate_stored
    #print 'R-avail', Node.basetime/(repair_rate_avail*60)
    #print 'R-stored', Node.basetime/(repair_rate_stored*60)
    #print f
    #return

    sim = Simulation()
    sim.initialize()
    if type=='avail':
        sto = StorageSystem(n*f, repair_rate_avail, threshold, required, sim=sim)
        #print sto.targeted_nodes
        #return
    elif type=='stored':
        sto = StorageSystem(n, repair_rate_stored, threshold, required, sim=sim)
        #print sto.targeted_nodes
        #return
    else:
        print 'error'
        return

    sim.activate(sto, sto.run())
    sim.simulate(until=100000000*Node.basetime)

if __name__=='__main__':
    if main.func_code.co_argcount != len(sys.argv)-1:
        print 'incorrect number of arguments'
    else:
        try:
            main(*sys.argv[1:])
        except KeyboardInterrupt:
            pass
